/**
 * 分割成>,=,<
 * @param {number[]} arr
 * @param {number} L
 * @param {number} R
 * @param {number} pivot
 */
import { swap } from '../sort/utils'

function partition (arr, L, R, pivot) {
  let less = R + 1
  let more = L - 1
  while (L < less) {
    if (arr[L] > pivot) {
      swap(arr, L++, ++more)
    } else if (arr[L] < pivot) {
      swap(arr, L, --less)
    } else {
      L++
    }
  }
  return [more + 1, less - 1]
}

/**
 * 插入排序
 * @param {number[]} arr
 * @param {number} L
 * @param {number} R
 */
function insertSort (arr, L, R) {
  for (let i = L + 1; i <= R; i++) {
    for (let j = i; j > L; j--) {
      if (arr[j] < arr[j - 1]) {
        swap(arr, j, j - 1)
      } else {
        break
      }
    }
  }
}

/**
 * 获取组内中位数
 * @param {number[]} arr
 * @param {number} L
 * @param {number} R
 * @return {number}
 */
function getMedian (arr, L, R) {
  insertSort(arr, L, R)
  const median = L + Math.floor((R - L) / 2)
  return arr[median]
}

/**
 * 获取组之间中位数
 * @param {number[]} arr
 * @param {number} L
 * @param {number} R
 */
function getMedianOfMedians (arr, L, R) {
  const counts = R - L + 1
  const offset = counts % 5 === 0 ? 0 : 1 // 可能有一组不够5个数
  const medianArr = new Array(Math.floor(counts / 5) + offset)
  for (let i = 0; i < medianArr.length; i++) {
    const begin = i * 5 + L
    const end = begin + 4
    medianArr[i] = getMedian(arr, begin, Math.min(end, R))
  }
  return select(medianArr, 0, medianArr.length - 1, Math.floor(medianArr.length / 2))
}

/**
 * 找第K大的数(K从0开始)
 * @param {number[]} arr
 * @param {number} L
 * @param {number} R
 * @param {number} K
 * @return {number}
 */
function select (arr, L, R, K) {
  if (L === R) {
    return arr[L]
  }
  const pivot = getMedianOfMedians(arr, L, R)
  const [more, less] = partition(arr, L, R, pivot)
  if (more <= K && less >= K) {
    return arr[K]
  } else if (more > K) {
    return select(arr, L, more - 1, K)
  } else {
    return select(arr, less + 1, R, K)
  }
}

/**
 * 获取top K元素，1<=K<=N
 * @param {number[]} arr
 * @param {number} K
 * @return {number}
 */
function topK (arr, K) {
  return select(arr, 0, arr.length - 1, K - 1)
}

console.log(topK([7, 6, 5, 4, 3, 2, 1], 5))
console.log(topK([99, 99], 1))
